import MySQLdb
import Util
import string, datetime, sets

INSERT_BUFFER_SIZE = 10000 #number of inserts to buffer before inserting
LOG_TABLE_CREATION_STRING = "CREATE TABLE IF NOT EXISTS logs (monitor ENUM('route-views', 'route-views2', 'route-views3',\
 'route-views4', 'route-views5', 'route-views6', 'rrc00', 'rrc01', 'rrc02', 'rrc03', 'rrc04', 'rrc05', 'rrc06', 'rrc07',\
 'rrc08', 'rrc09', 'rrc10', 'rrc11', 'rrc12') NOT NULL,\
    start_date DATETIME NOT NULL,\
    end_date DATETIME NOT NULL,\
    PRIMARY KEY(monitor),\
    INDEX(monitor))"
    
MASTER_TABLE_CREATION_STRING = "CREATE TABLE IF NOT EXISTS master (id INT UNSIGNED AUTO_INCREMENT,\
    ip INT UNSIGNED NOT NULL,\
    mask TINYINT UNSIGNED NOT NULL,\
    monitor ENUM('route-views', 'route-views2', 'route-views3', 'route-views4', 'route-views5', 'route-views6',\
 'rrc00', 'rrc01', 'rrc02', 'rrc03', 'rrc04', 'rrc05', 'rrc06', 'rrc07', 'rrc08', 'rrc09', 'rrc10', 'rrc11', 'rrc12') NOT NULL,\
    PRIMARY KEY(id),\
    UNIQUE INDEX(ip, mask, monitor))"

PLAIN_UPDATES_TABLE_CREATION_STRING = "CREATE TABLE IF NOT EXISTS %s (prefix_id INT UNSIGNED NOT NULL,\
    time DATETIME NOT NULL,\
    type ENUM('A','W') NOT NULL,\
    peer_ip INT UNSIGNED NOT NULL,\
    peer_as SMALLINT UNSIGNED NOT NULL,\
    as_path VARCHAR(120))"
    
UPDATE_TABLE_STR = "%s_%s_updates"
    
class UpdateDBHelper:
    def __init__(self, table, monitor=None):
        self.connection = MySQLdb.connect(host = "localhost",
                                          user = "user",
                                          passwd = "passwd",
                                          db = "updates")
        self.cursor = self.connection.cursor()

        self.cursor.execute(MASTER_TABLE_CREATION_STRING)
        self.cursor.execute(LOG_TABLE_CREATION_STRING)
        self.cursor.execute("SET time_zone='+00:00';")
        self.cursor.execute("SET bulk_insert_buffer_size=268435456;") #set bulk_insert_buffer_size to 256MB
        self.tableHack = table
        self.tableHackIndex = 0
        self.connection.commit()
        self.monitor = monitor
        self.curTable = None
        self.tablesUsed = sets.Set()
        self.insertBuffer = [] 
        
    def checkDates(self, startDate, endDate):
        self.cursor.execute("SELECT start_date, end_date FROM logs WHERE monitor='%s'" % self.monitor)
        dates = self.cursor.fetchone()
        if dates:
            assert((endDate + datetime.timedelta(seconds=1) == dates[0]) or (startDate == dates[1] + datetime.timedelta(seconds=1)))
            
    def setTable(self, month, year):
        self.flushUpdates()
        self.curTable = UPDATE_TABLE_STR % (year, month)
        self.tablesUsed.add(self.curTable)
        self.cursor.execute(PLAIN_UPDATES_TABLE_CREATION_STRING % (self.curTable))
        indexName = "%s_index" % self.curTable
        try:
            self.cursor.execute("DROP INDEX %s ON %s" % indexName, self.curTable)
        except:
            print "Index doesn't exist: %s" % indexName
        self.connection.commit()

    def createIndices(self):
        for table in self.tablesUsed:
            indexName = "%s_index" % table
            try:
                self.cursor.execute("CREATE INDEX %s ON %s (prefix_id)" % (indexName, table))
            except:
                print "Index on %s already exists" % table
        
    def getPrefixId(self, prefix):
        assert(self.monitor)
        (dotIp, mask) = prefix.split("/")
        
        self.cursor.execute("INSERT IGNORE INTO master VALUES(NULL, INET_ATON('%s'), %s, '%s')" %\
                            (dotIp,
                             mask,
                             self.monitor))
        
        self.cursor.execute("SELECT id FROM master WHERE ip=INET_ATON('%s') AND mask=%s" % (dotIp, mask))
        return self.cursor.fetchone()[0]
    
    def getMasterEntry(self, prefix):
        (dotIp, mask) = prefix.split("/")
        self.cursor.execute("SELECT id, ip, mask FROM master WHERE ip=INET_ATON('%s') AND mask=%s" % (dotIp, mask))
        return self.cursor.fetchone()
            
    def addUpdate(self, id, update):
        if update.as_path:
            asPathStr = "'%s'" % update.as_path
        else:
            asPathStr = 'NULL'
                   
        self.insertBuffer.append("(%s, FROM_UNIXTIME(%s), '%s', INET_ATON('%s'), %s, %s)" %\
                            (id,
                             update.time,
                             update.type,
                             update.peer_ip,
                             update.peer_as,
                             asPathStr))
        
        if (len(self.insertBuffer) == INSERT_BUFFER_SIZE):
            self.cursor.execute("INSERT INTO %s VALUES %s" % (self.curTable, ','.join(self.insertBuffer)))
            self.insertBuffer = []
        
    def setDates(self, startDate, endDate):
        assert(self.monitor)
        self.cursor.execute("INSERT INTO logs VALUES('%s', '%s', '%s')\
        ON DUPLICATE KEY UPDATE start_date=IF(start_date > '%s', '%s', start_date), end_date=IF(end_date < '%s', '%s', end_date)" %\
        (self.monitor, startDate, endDate, startDate, startDate, endDate, endDate))
        
    def getPrefixList(self):
        self.cursor.execute("SELECT id, ip, mask FROM master")
        return self.cursor.fetchall()
    
    def loadUpdatesHack(self):
        #AND time >= '%s' AND time <= '%s'
        self.cursor.execute("SELECT time, type, peer_as, as_path FROM %s WHERE prefix_id=%s ORDER BY time ASC, type DESC;" % (self.tableHack, self.id))
        
    def loadUpdates(self, id, startDate, endDate, monitor='route-views2'):
        self.id = id
        self.startDate = startDate
        self.endDate = endDate
        self.loadUpdatesHack()
        
    def fetchUpdates(self, quantity=200):
        res = self.cursor.fetchmany(quantity)
        return res
        
    def flushUpdates(self):
        if (len(self.insertBuffer) > 0):
            self.cursor.execute("INSERT INTO %s VALUES %s" % (self.curTable, ','.join(self.insertBuffer)))
            self.insertBuffer = [] 
            self.connection.commit()
                    
    def commit(self):       
        self.connection.commit()